import asyncio
import websockets
import argparse
import os
import time
from http.server import SimpleHTTPRequestHandler
from socketserver import TCPServer
from pathlib import Path

# ---------------------- WebSocket 服务器（接收并广播图片） ----------------------
clients = set()  # 保存所有连接的 WebSocket 客户端（前端）


async def handle_websocket(websocket, path):
    clients.add(websocket)
    try:
        async for message in websocket:
            # 1. 保存图片到本地
            timestamp = int(time.time())
            filename = f"image_{timestamp}.jpg"
            save_path = "./received_image/" + filename

            with open(save_path, "wb") as f:
                f.write(message)
            print(f"图片已保存至: {save_path}")

            # 2. 广播给所有前端客户端
            await broadcast_image(message)

    except websockets.exceptions.ConnectionClosed:
        print("客户端连接关闭")
    finally:
        clients.remove(websocket)


async def broadcast_image(image_data):
    # 向所有前端客户端发送图片数据
    if clients:
        await asyncio.gather(*[client.send(image_data) for client in clients])


# ---------------------- HTTP 静态服务器（托管前端页面） ----------------------
def run_http_server(port=8000):
    handler = SimpleHTTPRequestHandler
    httpd = TCPServer(("", port), handler)
    print(f"HTTP 服务器运行在 http://localhost:{port}")
    httpd.serve_forever()


# ---------------------- 主逻辑 ----------------------
async def main():
    parser = argparse.ArgumentParser(description="WebSocket 图片服务器")
    parser.add_argument("--port", type=int, default=8765, help="WebSocket 端口")
    parser.add_argument("--http_port", type=int, default=8000, help="HTTP 服务器端口")
    parser.add_argument("--save_dir", type=str, default="received_images", help="保存目录")
    args = parser.parse_args()

    # 创建保存目录
    Path(args.save_dir).mkdir(parents=True, exist_ok=True)

    # 启动 HTTP 服务器（在独立线程中）
    import threading
    http_thread = threading.Thread(target=run_http_server, args=(args.http_port,))
    http_thread.daemon = True
    http_thread.start()

    # 启动 WebSocket 服务器
    async with websockets.serve(handle_websocket, "0.0.0.0", args.port):
        print(f"WebSocket 服务器运行在 ws://0.0.0.0:{args.port}")
        await asyncio.Future()  # 永久运行


if __name__ == "__main__":
    asyncio.run(main())